#!/usr/bin/env python3
"""
Simplified Population ADAM PBPK Modeling Demo

A simplified version for running population ADAM PBPK modeling with 100 individuals.
Generates comprehensive charts and data analysis.

Usage:
    python3 demo_population_adam.py

Author: PBPK Modeling System
Version: 1.0.0
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os
import sys
from typing import Dict, List

# Add path
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))

from src.absorption.adam_model import ADAMModel
from src.absorption.absorption_module import AbsorptionModel
from src.population.population_data import PopulationData

def run_population_adam_modeling():
    """Run simplified population ADAM modeling for 100 individuals"""
    print("🧬 Population ADAM PBPK Modeling Demo")
    print("=" * 50)
    
    # Setup
    population_size = 100
    population_data = PopulationData()
    
    # Generate population
    print("📊 Generating population...")
    individuals = population_data.get_population_parameters('adult', population_size)
    
    # Drug parameters
    base_drug = {
        'drug_name': 'Test_Drug',
        'dose': 100,
        'solubility': 0.8,
        'permeability': 2e-4,
        'molecular_weight': 325.4,
        'pKa': 7.8,
        'logP': 2.5,
        'particle_size': 50,
        'density': 1.2,
        'use_adam': True
    }
    
    # Run simulations
    print("🔬 Running ADAM simulations...")
    time_points = np.linspace(0, 24, 100)
    results = []
    
    for i, individual in enumerate(individuals):
        # Create individual drug parameters
        drug_params = base_drug.copy()
        weight = individual.get('weight', 70)
        
        drug_params.update({
            'id': f'individual_{i+1:03d}',
            'weight': weight,
            'age': individual.get('age', 35),
            'dose': 100 * (weight / 70) ** 0.75,  # Allometric scaling
            'solubility': np.random.normal(0.8, 0.2),
            'permeability': np.random.lognormal(np.log(2e-4), 0.5),
            'pKa': np.random.normal(7.8, 0.5)
        })
        
        # Ensure bounds
        drug_params['solubility'] = max(0.1, min(2.0, drug_params['solubility']))
        drug_params['permeability'] = max(1e-5, min(1e-3, drug_params['permeability']))
        drug_params['pKa'] = max(5.0, min(10.0, drug_params['pKa']))
        
        # Run ADAM model
        adam_model = ADAMModel(drug_params)
        
        # BCS classification using AbsorptionModel
        absorption_model = AbsorptionModel(drug_params)
        bcs = absorption_model.calculate_bcs_classification(
            drug_params['solubility'],
            drug_params['permeability'],
            drug_params['dose']
        )
        
        # ADAM absorption
        adam_result = adam_model.calculate_adam_absorption(time_points)
        
        # Simple PK model
        bioavailability = adam_result['bioavailability']
        dose = drug_params['dose'] * bioavailability
        ke = 0.15 * (70 / weight) ** 0.25
        Vd = 50 * (weight / 70)
        ka = 1.0
        
        # Concentration calculation
        if abs(ka - ke) > 1e-6:
            conc = (dose * ka) / (Vd * (ka - ke)) * (
                np.exp(-ke * time_points) - np.exp(-ka * time_points)
            )
        else:
            conc = (dose * ka * time_points * np.exp(-ke * time_points)) / Vd
        
        conc = np.maximum(conc, 0)
        
        # PK metrics
        cmax = np.max(conc)
        tmax = time_points[np.argmax(conc)]
        try:
            from numpy import trapezoid
            auc = trapezoid(conc, time_points)
        except ImportError:
            # Fallback for older numpy versions
            auc = np.trapz(conc, time_points)
        
        results.append({
            'id': drug_params['id'],
            'age': individual.get('age', 35),
            'weight': weight,
            'bcs_class': bcs['class'],
            'solubility': drug_params['solubility'],
            'permeability': drug_params['permeability'],
            'cmax': cmax,
            'tmax': tmax,
            'auc': auc,
            'bioavailability': bioavailability,
            'concentration': conc,
            'time_points': time_points
        })
    
    # Generate visualizations
    print("📈 Generating charts...")
    generate_charts(results)
    
    # Export data
    print("💾 Exporting data...")
    export_data(results)
    
    # Summary
    print("\n📋 Population Summary:")
    print("-" * 30)
    bcs_counts = {cls: sum(1 for r in results if r['bcs_class'] == cls) 
                  for cls in ['I', 'II', 'III', 'IV']}
    for cls, count in bcs_counts.items():
        print(f"BCS Class {cls}: {count} individuals ({count/len(results)*100:.1f}%)")
    
    cmax_values = [r['cmax'] for r in results]
    auc_values = [r['auc'] for r in results]
    
    print(f"\nPK Parameters:")
    print(f"Cmax: {np.mean(cmax_values):.2f} ± {np.std(cmax_values):.2f} mg/L")
    print(f"AUC: {np.mean(auc_values):.1f} ± {np.std(auc_values):.1f} mg·h/L")
    print(f"CV Cmax: {np.std(cmax_values)/np.mean(cmax_values)*100:.1f}%")
    print(f"CV AUC: {np.std(auc_values)/np.mean(auc_values)*100:.1f}%")
    
    print("\n✅ Complete! Check output folder for results.")

def generate_charts(results: List[Dict]):
    """Generate population charts"""
    os.makedirs('../output', exist_ok=True)
    
    time_points = results[0]['time_points']
    
    # Create comprehensive visualization
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    fig.suptitle('Population ADAM PBPK Results (100 Individuals)', fontsize=16)
    
    # 1. Population concentration profile
    ax1 = axes[0, 0]
    concentrations = np.array([r['concentration'] for r in results])
    mean_conc = np.mean(concentrations, axis=0)
    std_conc = np.std(concentrations, axis=0)
    
    ax1.plot(time_points, mean_conc, 'b-', linewidth=2, label='Mean')
    ax1.fill_between(time_points, mean_conc-std_conc, mean_conc+std_conc, 
                     alpha=0.3, color='blue', label='±1 SD')
    ax1.set_xlabel('Time (hours)')
    ax1.set_ylabel('Concentration (mg/L)')
    ax1.set_title('Population Concentration Profile')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # 2. Individual trajectories (sample)
    ax2 = axes[0, 1]
    for i in range(0, min(10, len(results))):
        ax2.plot(time_points, results[i]['concentration'], alpha=0.7, linewidth=1)
    ax2.plot(time_points, mean_conc, 'r-', linewidth=2, label='Mean')
    ax2.set_xlabel('Time (hours)')
    ax2.set_ylabel('Concentration (mg/L)')
    ax2.set_title('Individual Trajectories')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # 3. Cmax distribution
    ax3 = axes[0, 2]
    cmax_values = [r['cmax'] for r in results]
    ax3.hist(cmax_values, bins=20, alpha=0.7, color='skyblue', edgecolor='black')
    ax3.axvline(np.mean(cmax_values), color='red', linestyle='--', 
                label=f'Mean: {np.mean(cmax_values):.2f}')
    ax3.set_xlabel('Cmax (mg/L)')
    ax3.set_ylabel('Frequency')
    ax3.set_title('Cmax Distribution')
    ax3.legend()
    
    # 4. AUC distribution
    ax4 = axes[1, 0]
    auc_values = [r['auc'] for r in results]
    ax4.hist(auc_values, bins=20, alpha=0.7, color='lightgreen', edgecolor='black')
    ax4.axvline(np.mean(auc_values), color='red', linestyle='--', 
                label=f'Mean: {np.mean(auc_values):.1f}')
    ax4.set_xlabel('AUC (mg·h/L)')
    ax4.set_ylabel('Frequency')
    ax4.set_title('AUC Distribution')
    ax4.legend()
    
    # 5. BCS classification
    ax5 = axes[1, 1]
    bcs_counts = [sum(1 for r in results if r['bcs_class'] == cls) 
                  for cls in ['I', 'II', 'III', 'IV']]
    colors = ['#ff9999', '#66b3ff', '#99ff99', '#ffcc99']
    ax5.pie(bcs_counts, labels=['I', 'II', 'III', 'IV'], colors=colors, autopct='%1.1f%%')
    ax5.set_title('BCS Classification Distribution')
    
    # 6. Parameter correlations
    ax6 = axes[1, 2]
    weights = [r['weight'] for r in results]
    ax6.scatter(weights, cmax_values, alpha=0.6, color='purple')
    z = np.polyfit(weights, cmax_values, 1)
    p = np.poly1d(z)
    ax6.plot(weights, p(weights), "r--", alpha=0.8)
    ax6.set_xlabel('Weight (kg)')
    ax6.set_ylabel('Cmax (mg/L)')
    ax6.set_title('Weight vs Cmax')
    ax6.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('../output/demo_population_adam_results.png', dpi=300, bbox_inches='tight')
    plt.show()

def export_data(results: List[Dict]):
    """Export population data to CSV"""
    os.makedirs('../output', exist_ok=True)
    df = pd.DataFrame([
        {
            'id': r['id'],
            'age': r['age'],
            'weight': r['weight'],
            'bcs_class': r['bcs_class'],
            'solubility': r['solubility'],
            'permeability': r['permeability'],
            'cmax': r['cmax'],
            'tmax': r['tmax'],
            'auc': r['auc'],
            'bioavailability': r['bioavailability']
        }
        for r in results
    ])
    
    df.to_csv('../output/demo_population_adam_summary.csv', index=False)

if __name__ == "__main__":
    run_population_adam_modeling()